#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import pandas as pd

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

# from factor_selection.gp_factor_selection_v1 import get_gp_factor
from factor_selection.get_gp_factor import get_gp_factor

from base.GPModelSave import load_gpmodel

from base.KeyTransform import stock_code_map
from base.KeyTransform import trade_date_map

''''''
class GPFactorToLib(object):
    def __init__(self,
                 raw_data_dir,
                 raw_data_filename,
                 model_dir,
                 save_factor_db='gp_factors',
                 save_barra_stripping_factor_db='gp_factors/barra_stripping',
                 label_data_dir='./zbc_gplearn_factor_mining/label_data',
                 label_data_filename='label_data'):
        self.raw_data_dir = raw_data_dir
        self.raw_data_filename = raw_data_filename

        self.label_data_dir = label_data_dir
        self.label_data_filename = label_data_filename

        self.model_dir = model_dir

        self.save_factor_db = save_factor_db
        self.save_barra_stripping_factor_db = save_barra_stripping_factor_db

    def initialization(self, raw_data=None, concat_label_data=None):
        if raw_data is None:
            # TODO - 读取数据
            test_X_data = pd.read_hdf(os.path.join(self.raw_data_dir, self.raw_data_filename + '.h5'))

        if concat_label_data is None:
            concat_label_data = pd.read_hdf(os.path.join(self.label_data_dir, self.label_data_filename + '.h5'))

        test_X_data = test_X_data.reset_index()

        test_X_data = trade_date_map(test_X_data, key='date', reverse=False)
        test_X_data = stock_code_map(test_X_data, key='stock_code', reverse=False)

        self.test_X_data = test_X_data.set_index(['date', 'stock_code'])

        self.neu_keys = concat_label_data.loc[self.test_X_data.index, ['ci1_code', 'cap']].copy()
        self.metric_features_keys = self.test_X_data.index

    def run(self, model_name_list, ts_base=5):
        for model_name in model_name_list:
            model = load_gpmodel(os.path.join(self.model_dir, model_name + '.pkl'))

            # TODO - 生成因子保存到db里面
            factor_data, \
            bfactor_data = get_gp_factor(raw_data=self.test_X_data,
                                         keys=self.metric_features_keys,
                                         neu_keys=self.neu_keys,
                                         model=model,
                                         factor_name=model_name,
                                         by_name=True,
                                         key_by_path=False,
                                         ts_base=ts_base,
                                         key_reverse=True)

            if len(factor_data.index.shape) == 2:
                factor_data = factor_data.reset_index()
            else:
                factor_data = factor_data.reset_index(drop=True)

            if len(bfactor_data.index.shape) == 2:
                bfactor_data = bfactor_data.reset_index()
            else:
                bfactor_data = bfactor_data.reset_index(drop=True)

            # TODO - to factor lib
            data_reader = DataReader(db=self.save_factor_db)
            data_reader.update_data(factor_data)
            data_reader.create_factor_table(filename=model_name, main_columns=['date', 'stock_code'])

            data_reader = DataReader(db=self.save_barra_stripping_factor_db)
            data_reader.update_data(bfactor_data)
            data_reader.create_factor_table(filename=model_name, main_columns=['date', 'stock_code'])

            print(model_name, 'done!')


def gp_factor_to_lib(model_name_list,
                     model_dir = './zbc_gplearn_factor_mining/model/ths_user_behavior_factors',
                     test_data_dir = './zbc_gplearn_factor_mining/processed_X_data',
                     test_data_filename = 'processed_X_user_behavior_data_v1',
                     label_data_dir = './zbc_gplearn_factor_mining/label_data',
                     label_data_filename = 'label_data',
                     ts_base=5,
                     save_factor_db = 'gp_factors',
                     save_bfactor_db = 'gp_factors/barra_stripping'):
    # TODO - 读取数据
    # TODO - 测试
    test_X_data = pd.read_hdf(os.path.join(test_data_dir, test_data_filename + '.h5'))
    concat_label_data = pd.read_hdf(os.path.join(label_data_dir, label_data_filename + '.h5'))

    test_X_data = test_X_data.reset_index()
    test_X_data = trade_date_map(test_X_data, key='date', reverse=False)
    test_X_data = stock_code_map(test_X_data, key='stock_code', reverse=False)
    test_X_data = test_X_data.set_index(['date', 'stock_code'])

    neu_keys = concat_label_data.loc[test_X_data.index, ['ci1_code', 'cap']].copy()
    metric_features_keys = test_X_data.index

    # model_name_list = [
    #     'zbc_ths_user_hehavior_factor_mining_v2_v%d' % id for id in range(1, 41) if id not in [23, 27]
    # ]

    for model_name in model_name_list:
        # model_name = 'zbc_ths_user_hehavior_factor_mining_v2_v6'
        # model_name = 'zbc_ths_user_hehavior_factor_mining_v2_v12'
        model = load_gpmodel(os.path.join(model_dir, model_name + '.pkl'))

        # TODO - 生成因子保存到db里面
        factor_data, \
        bfactor_data = get_gp_factor(raw_data=test_X_data,
                                     keys=metric_features_keys,
                                     neu_keys=neu_keys,
                                     model=model,
                                     factor_name=model_name,
                                     by_name=True,
                                     key_by_path=False,
                                     ts_base=ts_base,
                                     key_reverse=True)

        if len(factor_data.index.shape) == 2:
            factor_data = factor_data.reset_index()
        else:
            factor_data = factor_data.reset_index(drop=True)

        if len(bfactor_data.index.shape) == 2:
            bfactor_data = bfactor_data.reset_index()
        else:
            bfactor_data = bfactor_data.reset_index(drop=True)

        # TODO - to factor lib
        data_reader = DataReader(db=save_factor_db)
        data_reader.update_data(factor_data)
        data_reader.create_factor_table(filename=model_name, main_columns=['date', 'stock_code'])

        data_reader = DataReader(db=save_bfactor_db)
        data_reader.update_data(bfactor_data)
        data_reader.create_factor_table(filename=model_name, main_columns=['date', 'stock_code'])

        print(model_name, 'done!')





